// STL
#include <fstream>                         // std::ofstream
#include <string>                          // std::string
#include <vector>                          // std::vector<>

// Boost
#include <boost/archive/text_oarchive.hpp> // boost::archive::text_oarchive

// jScience
#include "jScience/linalg.hpp"             // Matrix<>

// stats++
#include "statsxx/dataset.hpp"             // Matrix_to_DataSet(), shuffle_data_set(), partition_data_set()
#include "statsxx/machine_learning/NeuralNet.hpp" // NEURAL_NET


//
// DESC: Trains a MLP.
//
NEURAL_NET train_MLP(
                     NEURAL_NET            mlp,
                     // -----
                     const int             method,
                     //------
                     const int             nepoch_min,
                     const int             nepoch_max,
                     // -----
                     const double          lr,
                     const double          lr_min,
                     const double          lr_max,
                     // -----
                     const double          momentum,
                     // -----
                     const double          weight_penalty,
                     // -----
                     const double          qrprop_u,
                     const double          qrprop_d,
                     // -----
                     const double          scg_lambda,
                     const double          scg_sigma,
                     const double          scg_convg_iterfrac,
                     const double          scg_rk_tol,
                     // -----
                     const Matrix<double> &X_in,
                     const Matrix<double> &X_out,
                     // -----
                     const int             npts_per_batch,
                     // -----
                     const std::string     mlp_file
                     )
{
    const bool silent  = false;

    //=========================================================

    //=========================================================
    // INITIALIZATION
    //=========================================================

    DataSet dataset = Matrix_to_DataSet(
                                        X_in,
                                        X_out
                                        );

    shuffle_data_set(dataset);

    std::vector<DataSet> datasets = partition_data_set(
                                                       dataset
                                                       );

    //=========================================================
    // TRAIN
    //=========================================================

    mlp.train(
              method,
              //------
              nepoch_min,
              nepoch_max,
              // -----
              lr,
              lr_min,
              lr_max,
              // -----
              momentum,
              // -----
              weight_penalty,
              // -----
              qrprop_u,
              qrprop_d,
              // -----
              scg_lambda,
              scg_sigma,
              scg_convg_iterfrac,
              scg_rk_tol,
              // -----
              datasets[0],
              datasets[1],
              datasets[2],
              // -----
              npts_per_batch,
              // -----
              silent
              );

    //=========================================================
    // OUTPUT
    //=========================================================

    // SAVE MLP TO ARCHIVE
    {
        std::ofstream ofs(mlp_file, std::ios::out);

        boost::archive::text_oarchive oa(ofs);

        oa << mlp;
    }

    return mlp;
}
